from abc import ABC, abstractmethod

import torch

from ..model.model import NNModel
from ..utils.config import Configer


class Runner(ABC):
    """
    Runner基类，管理神经网络模型开发过程中的最基本资源，包括模型、配置、CUDA设备等
    """

    DEFAULT_RES_DIR = "./resources"
    LARGE_INT = 99999

    def __init__(
        self,
        model: NNModel,
        configer: Configer,
        resource_dir: str = None,
        gpu_id: int = -1,
    ) -> None:
        super().__init__()
        self.configer = configer
        self.cfg = configer.params
        self.gpu_id = gpu_id
        self.resource_dir = resource_dir if resource_dir is not None else self.DEFAULT_RES_DIR
        self.device = self._init_device(gpu_id)
        self.model = model.to(self.device)
        print("Using Device:", self.device)

    @abstractmethod
    def run(self):
        pass

    def _init_device(self, gpu_id):
        use_gpu = (
            gpu_id >= 0 and gpu_id < torch.cuda.device_count() if gpu_id is not None else False
        )
        return torch.device("cuda:" + str(gpu_id) if use_gpu else "cpu")


def to_device(obj, device):
    if isinstance(obj, (tuple, list)):
        obj = [to_device(x, device) for x in obj]
    elif isinstance(obj, dict):
        new_obj = dict()
        for k, v in obj.items():
            new_obj[k] = to_device(v, device)
        obj = new_obj
    elif torch.is_tensor(obj):
        obj = obj.to(device)
    return obj
